#!/usr/bin/env bash

# Show Slurm user limits and current values
# Homepage: https://github.com/OleHolmNielsen/Slurm_tools/

#
# Command usage:
#
function usage()
{
        cat <<EOF
Usage: $0 [-u username [-A account] [-p partition] [-M cluster] [-l limit] [-s sub-limit] | -h ]
where:
        -u username: Print user <username> (Default is current user)
        -A accountname: Print only account <accountname>
        -p partition: Print only Slurm partition <partition>
        -M cluster: Print only cluster <cluster>
        -l Print selected limits only
        -s Print selected sub-limits only
        -h Print help information
EOF
}

# Default: Current user
username="$USER"

#
# Parse command options
#
while getopts "u:A:p:M:l:s:h" options; do
	case $options in
		u ) username=$OPTARG	# Select user
	    	;;
		A ) account=$OPTARG	# Select account
	    	;;
		p ) export partition=$OPTARG	# Select partition
	    	;;
		M ) export cluster=$OPTARG	# Select cluster
		    clusterarg="-M $OPTARG"
	    	;;
		l ) export limits=$OPTARG	# Select limits to display
	    	;;
		s ) export sublimits=$OPTARG	# Select sublimits to display
	    	;;
		h|*|? ) usage
	    	exit 1;;
	esac
done

# Check the username
if ! getent passwd $username > /dev/null
then
	echo ERROR: User $username does not exist in the passwd database
	exit 1
fi

# Check the account
if test -n "$account"
then
	if test -z "`sacctmgr -nrp show account $account`"
	then
		echo ERROR: Account $account does not exist in the Slurm database
		exit 1
	fi
else
	# Default account if none was specified
	account=`sacctmgr -nrp show user $username | awk -F'|' '{print $2}'`
fi

# Check the cluster name
if test -n "$cluster" -a -z "`sacctmgr -nrP show cluster $cluster format=cluster`"
then
	echo ERROR: Cluster $cluster does not exist in the Slurm database
	exit 1
fi

# Check the partition name
if test -n "$partition"
then
	# User $clusterarg to select a different cluster
	if ! scontrol -o $clusterarg show partition $partition >/dev/null
	then
		echo ERROR: Partition $partition does not exist in cluster $cluster
		exit 1
	fi
fi

# Test for extraneous command line arguments
if test $# -gt $(($OPTIND-1))
then
	echo ERROR: Too many command line arguments: $*
	usage
	exit 1
fi

# Get Current Association Manager state Association Records
# Valid assoc_mgr flags are 'Assoc, QOS, and/or Users'
scontrol -o show assoc_mgr users=$username account=$account flags=Assoc | awk '
BEGIN {
	partition	= ENVIRON["partition"]
	cluster		= ENVIRON["cluster"]
	limits		= ENVIRON["limits"]
	# Add some aliases for limits:
	if (limits == "AssocGrpCPURunMinutesLimit")
		limits = "GrpTRESRunMins"
	else if (limits == "AssocGrpCpuLimit")
		limits = "GrpTRESMins"
	else if (limits == "AssocGrpCpuLimit")
		limits = "GrpTRES"
	else if (limits == "AssocGrpGRES")
		limits = "GrpTRES"
	sublimits	= ENVIRON["sublimits"]
	# Add some aliases for sublimits:
	if (sublimits == "gpu")
		sublimits = "gres/gpu"
}
function split_parenthesis (string, array) {
	# Split strings of the form aaa(bbb) and return array
	n = split(string,array,"(")
	if (n <= 1) return 0	# No ( found
	if (array[1] == "N") array[1] = "None"		# Value N means None
	gsub(")","",array[2])	# Remove trailing )
	return 1
}
{
	# Read ClusterName lines
	split($1,a,"=")
	# Skip non-Association records
	if (a[1] != "ClusterName")	# ClusterName is field no. 1
		next
	# Selected cluster: Get the ClusterName record and skip if irrelevant
	if (cluster != "") {
		if (a[2] != "" && a[2] != cluster)
			next
	}

	# Selected partition: Get the Partition record and skip if irrelevant
	if (partition != "") {
		split($4,a,"=")	# Partition is field no. 4
		if (a[2] != "" && a[2] != partition)
			next
	}

	# Determine parent/user association records
	split($3,a,"=")	# UserName is field no. 3
	if (a[2] == "")	# Empty UserName field
		print "Association (Parent account):"
	else
		print "Association (User):"

	# Loop over association records
	for (i=1; i<=NF; i++) {	
		numrec = split($i,a,"=")
		rec = a[1]
		value = a[2]
		printok = 0
		if (limits == "" || rec == limits)
			printok = 1
		# Always print these lines
		if (rec == "ClusterName" || rec == "Account" || rec == "UserName" || rec == "Partition")
			printok = 1
		if (printok && numrec >= 2) {			# Print further fields
			if (rec == "UserName" && value == "")	# Empty UserName
				value = "None(Parent account)"
			else if (rec == "Partition" && value == "")
				value = "None(Any partition)"	# Empty partition = any partition
			printf("\t%14s = ", rec)		# Print record
			if (numrec == 2) {
				m = split_parenthesis(value, array)
				if (m == 0)
					printf("\t%s", value)	# Value
				else if (rec == "UserName" && array[1] != "None")
					printf("\t%s, UID=%s", array[1], array[2])	# UserName and UID
				else
					printf("\t%s, current value = %s", array[1], array[2])	# Values
			} else if (numrec >= 3) {			# Extra value fields
				printf("\n")
				gsub(rec "=","",$i)		# Remove leading rec= from string
				m = split($i,b,",")		# Comma-separated sub-fields
				for (j=1; j<=m; j++) {
					split(b[j],c,"=")	# Sub-limit value=field
					if (sublimits == "" || c[1] == sublimits) {
						# Value c[2] sub-field consists of "Limit(Value)" pairs
						if (split_parenthesis(c[2], array))
							printf("\t\t%8s:\tLimit = %s, current value = %s\n",
								c[1], array[1], array[2])
						else	# Just in case c[2] might not consist of "Limit(Value)"
							printf("\t\t%8s:\tLimit(Value) = %s\n", c[1], c[2])
					}
				}
			}
			printf("\n")
		}
	}
}'
